package com.xw.screen.capture.util;

import com.xw.screen.capture.entity.ScreenImageInfo;
import lombok.extern.slf4j.Slf4j;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.*;

/**
 * @author ：xiongwei
 * @date ：Created in 2020/11/21 23:34
 */
@Slf4j
public class CaptureScreenUtils {


    /**
     * 截取当前屏幕，返回图片名称
     *
     * @return 图片名称
     * @throws IOException  io操作异常
     * @throws AWTException 截屏异常
     */
    public static ScreenImageInfo createScreenImage() throws IOException, AWTException {
        return createScreenImage(1.0, 1.0);
    }

    /**
     * 截取当前屏幕，返回图片名称
     *
     * @return 图片名称
     * @throws IOException  io操作异常
     * @throws AWTException 截屏异常
     */
    public static ScreenImageInfo createScreenImageFour() throws IOException, AWTException {
        return createScreenImage(0.4, 0.4);
    }

    /**
     * 截取当前屏幕，返回图片名称
     *
     * @param widthProportion  截取宽度比例，取值0.0 ~ 1.0之间
     * @param heightProportion 高度比例，取值0.0 ~ 1.0之间
     * @return 图片名称
     * @throws IOException  io操作异常
     * @throws AWTException 截屏异常
     */
    public static ScreenImageInfo createScreenImage(Double widthProportion, Double heightProportion) throws IOException, AWTException {

        long time = System.currentTimeMillis();
        BufferedImage bufImage = screenImage(widthProportion, heightProportion);
        long time2 = System.currentTimeMillis();

        ScreenImageInfo screenImageInfo = getScreenImageInfo(time);
        ImageIO.write(bufImage, "PNG", new File(screenImageInfo.getPath() + File.separator + screenImageInfo.getFileName()));

        long time3 = System.currentTimeMillis();
        log.info("time = {}!", time);
        log.info("time2 = {}!", time2);
        log.info("time3 = {}!", time3);
        return screenImageInfo;
    }

    private static BufferedImage screenImage(Double widthProportion, Double heightProportion) throws AWTException {
        if (widthProportion < 0) {
            widthProportion = 0.0;
        }
        if (heightProportion < 0) {
            heightProportion = 0.0;
        }
        if (widthProportion > 1.0) {
            widthProportion = 1.0;
        }
        if (heightProportion > 1.0) {
            heightProportion = 1.0;
        }

        // 获取屏幕尺寸
        Dimension screenSize = Toolkit.getDefaultToolkit().getScreenSize();
        // 创建需要截取的矩形区域
        Rectangle rect = new Rectangle(0, 0, (int) (screenSize.width * widthProportion), (int) (screenSize.height * heightProportion));
        return new Robot().createScreenCapture(rect);
    }


    private static ScreenImageInfo getScreenImageInfo(long time) {
        // 保存截取的图片
        String name = time + ".png";
        ScreenImageInfo screenImageInfo = new ScreenImageInfo();
        screenImageInfo.setFileName(name);
        screenImageInfo.setPath(System.getProperty("user.dir") + File.separator + "image");
        return screenImageInfo;
    }

    /**
     * 截取当前屏幕，返回图片名称
     *
     * @param widthProportion  截取宽度比例，取值0.0 ~ 1.0之间
     * @param heightProportion 高度比例，取值0.0 ~ 1.0之间
     * @return 图片名称
     * @throws IOException  io操作异常
     * @throws AWTException 截屏异常
     */
    public static ScreenImageInfo createInputImage(Double widthProportion, Double heightProportion) throws IOException, AWTException {

        long time = System.currentTimeMillis();
        BufferedImage bufImage = screenImage(widthProportion, heightProportion);
        long time2 = System.currentTimeMillis();
        ScreenImageInfo screenImageInfo = getScreenImageInfo(time);

        InputStream inputStream = bufferedImageToInputStream(bufImage);
        bufImage.flush();
        screenImageInfo.setInputStream(inputStream);

        long time3 = System.currentTimeMillis();
        log.info("time = {}!", time);
        log.info("time2 = {}!", time2);
        log.info("time3 = {}!", time3);
        return screenImageInfo;
    }

    /**
     * 截取当前屏幕，返回图片名称
     *
     * @param widthProportion  截取宽度比例，取值0.0 ~ 1.0之间
     * @param heightProportion 高度比例，取值0.0 ~ 1.0之间
     * @return 图片名称
     * @throws IOException  io操作异常
     * @throws AWTException 截屏异常
     */
    public static ScreenImageInfo createOutputImage(Double widthProportion, Double heightProportion) throws IOException, AWTException {

        long time = System.currentTimeMillis();
        BufferedImage bufImage = screenImage(widthProportion, heightProportion);
        long time2 = System.currentTimeMillis();
        ScreenImageInfo screenImageInfo = getScreenImageInfo(time);

        ByteArrayOutputStream os = new ByteArrayOutputStream();
        ImageIO.write(bufImage, "png", os);
        screenImageInfo.setOutputStream(os);

        long time3 = System.currentTimeMillis();
        log.info("time = {}!", time);
        log.info("time2 = {}!", time2);
        log.info("time3 = {}!", time3);
        return screenImageInfo;
    }

    /**
     * 将BufferedImage转换为InputStream
     *
     * @param image 截取屏幕的流
     * @return 输出流
     */
    public static ByteArrayInputStream bufferedImageToInputStream(BufferedImage image) {
        ByteArrayOutputStream os = new ByteArrayOutputStream();
        ByteArrayInputStream input = null;
        try {
            ImageIO.write(image, "png", os);
            input = new ByteArrayInputStream(os.toByteArray());
        } catch (IOException e) {
            log.error("提示:", e);
        } finally {
            try {
                os.close();
            } catch (IOException ignored) {
            }
        }
        return input;
    }

}
